import pandas
from typing import Optional
from utils.ai_translator.model import Model
from utils.ai_translator.translator.pdf_parser import PDFParser
from utils.ai_translator.translator.writer import Writer
from utils.ai_translator.translator_utils import LOG


class PDFTranslator:
    def __init__(self, model: Model, pdf_file_path: str, pages: Optional[int] = None, target_language: str = "中文"):
        self.model = model
        self.pdf_parser = PDFParser()
        self.writer = Writer()
        self.book = self.pdf_parser.parse_pdf(pdf_file_path, pages)
        self.target_language = target_language

    def get_book_language(self):
        all_content = ""
        for page_idx, page in enumerate(self.book.pages):
            for content_idx, content in enumerate(page.contents):
                if isinstance(content.original, str):
                    all_content += content.original
                elif isinstance(content.original, pandas.core.frame.DataFrame):
                    all_content += content.original.to_string()
                else:
                    pass

        prompt = self.model.make_language_prompt(all_content)
        translation, status = self.model.make_request(prompt)

        if status:
            return translation
        else:
            return None

    def translate_pdf(self,
                      file_format: str = 'MARKDOWN',
                      output_file_path: str = None):

        for page_idx, page in enumerate(self.book.pages):
            for content_idx, content in enumerate(page.contents):
                prompt = self.model.translate_prompt(content, self.target_language)
                LOG.debug(prompt)
                translation, status = self.model.make_request(prompt)
                LOG.info(translation)

                # Update the content in self.book.pages directly
                self.book.pages[page_idx].contents[content_idx].set_translation(translation, status)

        self.writer.save_translated_book(self.book, output_file_path, file_format)
